import streamlit as st
import pandas as pd
import pickle
import os
import json
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules
from collections import Counter
from zhipuai import ZhipuAI
import numpy as np
import re
from streamlit_echarts import st_echarts  # 新增ECharts库导入
import uuid
from io import BytesIO
import matplotlib.pyplot as plt

# 设置中文字体
# 初始化页面配置
st.set_page_config(page_title="招聘智能问答机器人", layout="wide", page_icon="🤖")

# 缓存配置
CACHE_ENABLED = True
FREQUENT_ITEMS_CACHE = "apriori_frequent_items.bin"
RULES_CACHE = "apriori_rules.bin"

# 初始化会话状态时添加charts存储
if "charts" not in st.session_state:
    st.session_state.charts = {}

# 数据预处理
@st.cache_data
def load_and_preprocess_data():
    try:
        # 这里修改为读取上传的文件
        data = pd.read_csv('拉勾网2223招聘数据.csv')
    except Exception as e:
        st.error(f"数据加载失败: {str(e)}")
        return pd.DataFrame()

    # 处理缺失值，简单删除含有缺失值的行
    data = data.dropna()

    # 处理异常值，使用四分位数间距识别和处理异常薪资数据
    Q1 = data['avg_salary'].quantile(0.25)
    Q3 = data['avg_salary'].quantile(0.75)
    IQR = Q3 - Q1
    lower_bound = Q1 - 1.5 * IQR
    upper_bound = Q3 + 1.5 * IQR
    data = data[(data['avg_salary'] >= lower_bound) & (data['avg_salary'] <= upper_bound)]

    # 清洗薪资数据
    data['avg_salary'] = data['avg_salary'].apply(clean_salary_data)

    # 技能列表预处理
    data['skill_list'] = data['skill_list'].apply(
        lambda x: [s.strip().lower() for s in str(x).split() if s.strip()]
        if pd.notnull(x) else []
    )

    data = data[data['skill_list'].apply(len) > 0]
    return data

# 薪资数据清洗函数
def clean_salary_data(salary):
    if isinstance(salary, str):
        # 处理“Xk - Yk”格式
        match = re.search(r'(\d+)k\s*-\s*(\d+)k', salary)
        if match:
            min_salary = int(match.group(1))
            max_salary = int(match.group(2))
            return (min_salary + max_salary) / 2
        # 处理“Xk以上”格式
        match = re.search(r'(\d+)k以上', salary)
        if match:
            return int(match.group(1))
        # 处理“X - Y万”格式
        match = re.search(r'(\d+)\s*-\s*(\d+)万', salary)
        if match:
            min_salary = int(match.group(1)) * 10
            max_salary = int(match.group(2)) * 10
            return (min_salary + max_salary) / 2
        # 处理“X万以上”格式
        match = re.search(r'(\d+)万以上', salary)
        if match:
            return int(match.group(1)) * 10
    return salary

# 关联规则生成
@st.cache_resource
def generate_association_rules(data):
    if CACHE_ENABLED and os.path.exists(FREQUENT_ITEMS_CACHE) and os.path.exists(RULES_CACHE):
        with open(FREQUENT_ITEMS_CACHE, 'rb') as f:
            frequent_itemsets = pickle.load(f)
        with open(RULES_CACHE, 'rb') as f:
            rules = pickle.load(f)
        return frequent_itemsets, rules

    te = TransactionEncoder()
    te_ary = te.fit_transform(data['skill_list'])
    df = pd.DataFrame(te_ary, columns=te.columns_)

    frequent_itemsets = apriori(df, min_support=0.02, use_colnames=True, max_len=3)

    if frequent_itemsets.empty:
        return pd.DataFrame(columns=['itemsets']), pd.DataFrame(columns=['antecedents', 'consequents'])

    try:
        rules = association_rules(frequent_itemsets, metric="lift", min_threshold=1.0)
    except ValueError:
        return frequent_itemsets, pd.DataFrame(columns=['antecedents', 'consequents'])

    valid_rules = rules[(rules['confidence'] >= 0.4) & (rules['lift'] > 1.0)]

    if CACHE_ENABLED:
        with open(FREQUENT_ITEMS_CACHE, 'wb') as f:
            pickle.dump(frequent_itemsets, f)
        with open(RULES_CACHE, 'wb') as f:
            pickle.dump(valid_rules, f)

    return frequent_itemsets, valid_rules

# 初始化AI客户端
def init_ai(api_key):
    return ZhipuAI(api_key=api_key)

# LLM交互封装
def chat_tools(client, prompt, tools=None, temperature=0.6):
    if tools:
        response = client.chat.completions.create(
            model="glm-4",
            messages=[{"role": "user", "content": prompt}],
            tools=tools,
            tool_choice="auto",
            temperature=temperature
        )
    else:
        response = client.chat.completions.create(
            model="glm-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=temperature
        )
    return response

# 自然语言生成模板
def generate_prompt(question, result, function_name, skill_name=None, position=None, city=None, company=None, education=None):
    if not result or (isinstance(result, str) and (result.startswith("数据集未收录") or result.startswith("没有找到"))):
        return f"""
        用户提问：{question}
        查询结果：在私有数据中未查到相关信息。
        请调用共有数据，基于行业普遍情况为用户提供相关回答。
        """
    else:
        if function_name == "recommend_related_skills":
            skill_name = skill_name or "未知技能"
            result_str = ", ".join([r[0] for r in result]) if isinstance(result[0], tuple) else ", ".join(result)
            return f"""
            用户提问：{question}
            查询结果：从私有数据中找到了'{skill_name}'技术的关联规则，关联技术为：{result_str}。
            请按以下格式回复用户：
            1. 根据您的技能，我在数据中查到了相关技术的信息，分别为：{result_str}。
            2. 简要说明这些技术与您技能的关联性和应用场景。
            """
        elif function_name == "query_skills_for_position":
            position = position or "未知岗位"
            result_str = ", ".join([f"{r[0]}（{r[1]}次）" for r in result]) if isinstance(result[0], tuple) else ", ".join(result)
            return f"""
            用户提问：{question}
            查询结果：'{position}'岗位通常需要的技能为：{result_str}。
            请按以下格式回复用户：
            1. 根据您的问题，在私有数据中我查到了与'{position}'岗位相关的技能需求信息。'{position}'岗位通常需要以下技能：{result_str}。
            2. 简要说明这些技能的重要性和应用场景。
            """
        elif function_name == "recommend_positions_by_skill":
            skill_name = skill_name or "未知技能"
            position_counts = Counter(result)
            position_list = [pos for pos, _ in position_counts.most_common()]
            position_str = ", ".join(position_list)

            return f"""
            用户提问：{question}
            查询结果：根据您会的技术，我在数据中查到了相关岗位的信息，分别为：{position_str}。
            请按以下格式回复用户：
            1. 根据您会的技术，我在数据中查到了相关岗位的信息，分别为：{position_str}。
            2. 简要说明每个岗位的主要职责和发展前景。
            """

        elif function_name == "query_highest_salary_jobs_by_city":
            city = city or "未知城市"
            if isinstance(result, str):
                return f"""
                用户提问：{question}
                查询结果：{result}。
                请以自然语言的方式，简洁易懂地进行回答。
                """
            else:
                job_info = []
                for idx, job in enumerate(result, start=1):
                    job_info.append(f"岗位{idx}（{job['position']}）薪资{job['salary']}元")
                job_info_str = ", ".join(job_info)
                return f"""
                用户提问：{question}
                查询结果：'{city}'薪资最高的五个岗位为：{job_info_str}。
                请以自然语言的方式，详细列出每个岗位及其对应的平均薪资，并简要说明这些岗位的发展前景。
                """

        elif function_name == "query_jobs_by_education":
            education = education or "未知学历"
            if isinstance(result, str):
                return f"""
                用户提问：{question}
                查询结果：{result}。
                请以自然语言的方式，简洁易懂地进行回答。
                """
            else:
                positions = ", ".join([job["positionName"] for job in result])
                return f"""
                用户提问：{question}
                查询结果：适合'{education}'学历的职业有：{positions}。
                请以自然语言的方式，详细说明这些职业的发展前景和薪资水平。
                """
        else:
            result_str = str(result)
            return f"""
            用户提问：{question}
            查询结果：{result_str}
            请以专业的语言回答用户的问题。
            """


# 自定义函数工具
recommend_related_skills_tool = {
    "type": "function",
    "function": {
        "name": "recommend_related_skills",
        "description": "根据输入技术推荐关联技术",
        "parameters": {
            "type": "object",
            "properties": {"skill": {"type": "string"}},
            "required": ["skill"]
        }
    }
}

query_skills_for_position_tool = {
    "type": "function",
    "function": {
        "name": "query_skills_for_position",
        "description": "根据岗位名称查询所需技术",
        "parameters": {
            "type": "object",
            "properties": {"position": {"type": "string"}},
            "required": ["position"]
        }
    }
}

recommend_positions_by_skill_tool = {
    "type": "function",
    "function": {
        "name": "recommend_positions_by_skill",
        "description": "根据技术推荐适合岗位",
        "parameters": {
            "type": "object",
            "properties": {"skill": {"type": "string"}},
            "required": ["skill"]
        }
    }
}


query_highest_salary_jobs_by_city_tool = {
    "type": "function",
    "function": {
        "name": "query_highest_salary_jobs_by_city",
        "description": "根据输入的城市名称，查询该城市薪资最高的岗位信息",
        "parameters": {
            "type": "object",
            "properties": {
                "city": {
                    "type": "string",
                    "description": "需要查询的城市名称"
                }
            },
            "required": ["city"]
        }
    }
}



query_jobs_by_education_tool = {
    "type": "function",
    "function": {
        "name": "query_jobs_by_education",
        "description": "根据输入的学历，查询适合该学历的五个职业",
        "parameters": {
            "type": "object",
            "properties": {
                "education": {
                    "type": "string",
                    "description": "需要查询的学历"
                }
            },
            "required": ["education"]
        }
    }
}

# 函数实现
def recommend_related_skills(skill, rules, data_skills):
    valid_skill = skill.strip().lower()
    if valid_skill not in data_skills:
        return [f"数据集未收录'{skill}'技术"]

    related = []
    for _, rule in rules.iterrows():
        antecedents = set(rule['antecedents'])
        if valid_skill in antecedents:
            related.extend([
                (tech, rule['confidence'], rule['lift']) 
                for tech in rule['consequents'] 
                if tech != valid_skill
            ])
    
    sorted_related = sorted(
        list({tech[0]: tech for tech in related}.values()),
        key=lambda x: -x[1]
    )
    
    return [(tech[0], tech[1]) for tech in sorted_related[:5]]  # 返回前5个关联技术

def query_skills_for_position(position, data):
    valid_position = position.strip().lower()
    if not data['positionName'].str.lower().str.contains(valid_position).any():
        return [f"数据集未收录'{position}'岗位"]
    
    relevant_rows = data[data['positionName'].str.lower().str.contains(valid_position)]
    all_skills = [skill for sublist in relevant_rows['skill_list'] for skill in sublist]
    
    skill_counts = Counter(all_skills)
    top_skills = skill_counts.most_common(5)
    
    return top_skills

def recommend_positions_by_skill(skill, data):
    valid_skill = skill.strip().lower()
    if not data['skill_list'].apply(lambda x: valid_skill in x).any():
        return [f"数据集未显示掌握'{skill}'可适配的岗位"]
    
    relevant_positions = data[data['skill_list'].apply(lambda x: valid_skill in x)]['positionName'].tolist()
    position_counts = Counter(relevant_positions)
    top_positions = position_counts.most_common(5)
    
    return top_positions

def query_salary_by_position(position, data):
    print(f"正在调用函数: query_salary_by_position")
    # 修复：增加城市过滤逻辑
    if " " in position:
        parts = position.split(" ")
        city = parts[0]
        position_name = " ".join(parts[1:])
        filtered_data = data[(data['positionName'] == position_name) & (data['city'] == city)]
        if filtered_data.empty:
            filtered_data = data[data['positionName'] == position_name]
    else:
        filtered_data = data[data['positionName'] == position]
        
    if filtered_data.empty:
        return "没有找到该岗位的相关数据"
    avg_salary = filtered_data['avg_salary'].mean() * 1000  # 转换为元
    return avg_salary

def count_jobs_by_city(city, data):
    print(f"正在调用函数: count_jobs_by_city")
    filtered_data = data[data['city'] == city]
    job_count = len(filtered_data)
    return job_count

def query_highest_salary_jobs_by_city(city, data):
    print(f"正在调用函数: query_highest_salary_jobs_by_city")
    filtered_data = data[data['city'] == city]
    if filtered_data.empty:
        return "没有找到该城市的相关岗位数据"
    # 按岗位分组计算平均薪资
    grouped = filtered_data.groupby('positionName')['avg_salary'].mean().reset_index()
    # 按平均薪资降序排序
    sorted_data = grouped.sort_values(by='avg_salary', ascending=False)
    # 取前五个岗位
    top_five_jobs = sorted_data.head(5)
    job_info = []
    for _, row in top_five_jobs.iterrows():
        job_info.append({
            "position": row['positionName'],
            "salary": int(row['avg_salary'] * 1000)  # 转换为元并取整
        })
    return job_info

def query_jobs_by_education(education, data):
    print(f"正在调用函数: query_jobs_by_education")
    filtered_data = data[data['education'] == education]
    if filtered_data.empty:
        return "没有找到适合该学历的岗位数据"
    
    # 按岗位分组并计算每个岗位的数量
    position_counts = filtered_data['positionName'].value_counts().reset_index()
    position_counts.columns = ['positionName', 'count']
    
    # 选择出现次数最多的5个岗位
    top_positions = position_counts.head(5)['positionName'].tolist()
    
    # 获取每个岗位的详细信息（职位名称、公司全称、城市、平均薪资、工作年限要求）
    result = []
    for position in top_positions:
        position_data = filtered_data[filtered_data['positionName'] == position].iloc[0]
        result.append({
            "positionName": position_data['positionName'],
            "companyFullName": position_data['companyFullName'],
            "city": position_data['city'],
            "avg_salary": int(position_data['avg_salary'] * 1000),  # 转换为元
            "workYear": position_data['workYear']
        })
    
    return result

# 将numpy类型转换为Python原生类型，确保ECharts能正确处理
def convert_numpy_types(data):
    if isinstance(data, np.ndarray):
        return data.tolist()
    elif isinstance(data, np.generic):
        return data.item()
    elif isinstance(data, dict):
        return {k: convert_numpy_types(v) for k, v in data.items()}
    elif isinstance(data, list):
        return [convert_numpy_types(v) for v in data]
    else:
        return data

# 使用ECharts的可视化函数
def visualize_related_skills(skill, related_skills):
    # 检查是否有有效的关联技能数据
    if not related_skills or isinstance(related_skills[0], str) and related_skills[0].startswith("数据集未收录"):
        return None
    
    skills = [skill[0] for skill in related_skills]
    confidence = [skill[1] for skill in related_skills]
    
    option = {
        "tooltip": {
            "trigger": 'axis',
            "axisPointer": {
                "type": 'shadow'
            },
            "formatter": "{b}: {c}"
        },
        "xAxis": {
            "type": "category",
            "data": skills,
            "axisLabel": {
                "rotate": 45
            }
        },
        "yAxis": {
            "type": "value",
            "min": 0,
            "max": 1,
            "axisLabel": {
                "formatter": "{value}"
            }
        },
        "series": [
            {
                "data": confidence,
                "type": "bar",
                "itemStyle": {
                    "color": "#4CAF50"
                }
            }
        ]
    }
    
    # 转换option中的numpy类型
    option = convert_numpy_types(option)
    
    # 生成唯一ID
    chart_id = str(uuid.uuid4())
    
    # 保存图表配置数据到会话状态
    st.session_state.charts[chart_id] = option
    
    # 保存图表数据到会话状态
    st_echarts(options=option, key=chart_id)
    
    return chart_id

def visualize_skills_for_position(position, skills):
    # 检查是否有有效的技能数据
    if not skills or isinstance(skills[0], str) and skills[0].startswith("数据集未收录"):
        return None
    
    skill_names = [skill[0] for skill in skills]
    counts = [skill[1] for skill in skills]
    
    option = {
        "tooltip": {
            "trigger": 'axis',
            "axisPointer": {
                "type": 'shadow'
            },
            "formatter": "{b}: {c}次"
        },
        "xAxis": {
            "type": "category",
            "data": skill_names,
            "axisLabel": {
                "rotate": 45
            }
        },
        "yAxis": {
            "type": "value"
        },
        "series": [
            {
                "data": counts,
                "type": "bar",
                "itemStyle": {
                    "color": "#2196F3"
                }
            }
        ]
    }
    
    # 转换option中的numpy类型
    option = convert_numpy_types(option)
    
    # 生成唯一ID
    chart_id = str(uuid.uuid4())
    
    # 保存图表配置数据到会话状态
    st.session_state.charts[chart_id] = option
    
    # 保存图表数据到会话状态
    st_echarts(options=option, key=chart_id)
    
    return chart_id

def visualize_positions_by_skill(skill, positions):
    # 检查是否有有效的岗位数据
    if not positions or isinstance(positions[0], str) and positions[0].startswith("数据集未显示"):
        return None
    
    position_names = [position[0] for position in positions]
    counts = [position[1] for position in positions]
    
    option = {
        "tooltip": {
            "trigger": 'axis',
            "axisPointer": {
                "type": 'shadow'
            },
            "formatter": "{b}: {c}次"
        },
        "xAxis": {
            "type": "category",
            "data": position_names,
            "axisLabel": {
                "rotate": 45
            }
        },
        "yAxis": {
            "type": "value"
        },
        "series": [
            {
                "data": counts,
                "type": "bar",
                "itemStyle": {
                    "color": "#FF9800"
                }
            }
        ]
    }
    
    # 转换option中的numpy类型
    option = convert_numpy_types(option)
    
    # 生成唯一ID
    chart_id = str(uuid.uuid4())
    
    # 保存图表配置数据到会话状态
    st.session_state.charts[chart_id] = option
    
    # 保存图表数据到会话状态
    st_echarts(options=option, key=chart_id)
    
    return chart_id

def visualize_highest_salary_jobs_by_city(city, jobs):
    # 检查是否有有效的高薪岗位数据
    if isinstance(jobs, str) or not jobs:
        return None
    
    job_names = [job['position'] for job in jobs]
    salaries = [job['salary'] for job in jobs]
    
    option = {
        "tooltip": {
            "trigger": 'axis',
            "axisPointer": {
                "type": 'shadow'
            },
            "formatter": "{b}: {c}元"
        },
        "xAxis": {
            "type": "category",
            "data": job_names,
            "axisLabel": {
                "rotate": 45
            }
        },
        "yAxis": {
            "type": "value",
            "axisLabel": {
                "formatter": "{value} 元"
            }
        },
        "series": [
            {
                "data": salaries,
                "type": "bar",
                "itemStyle": {
                    "color": "#FF5722"
                }
            }
        ]
    }
    
    # 转换option中的numpy类型
    option = convert_numpy_types(option)
    
    # 生成唯一ID
    chart_id = str(uuid.uuid4())
    
    # 保存图表配置数据到会话状态
    st.session_state.charts[chart_id] = option
    
    # 保存图表数据到会话状态
    st_echarts(options=option, key=chart_id)
    
    return chart_id

def visualize_jobs_by_education(education, jobs):
    # 检查是否有有效的职位数据
    if isinstance(jobs, str) or not jobs:
        return None
    
    # 创建一个DataFrame用于显示
    jobs_df = pd.DataFrame(jobs)
    
    # 重命名列以便更友好的显示
    jobs_df = jobs_df.rename(columns={
        'positionName': '职位名称',
        'companyFullName': '公司全称',
        'city': '城市',
        'avg_salary': '平均薪资(元)',
        'workYear': '工作年限要求'
    })
    
    # 使用Streamlit的表格显示
    st.subheader(f"适合'{education}'学历的职业")
    st.dataframe(jobs_df, use_container_width=True)
    
    # 生成唯一ID
    chart_id = str(uuid.uuid4())
    
    # 保存图表配置数据到会话状态
    st.session_state.charts[chart_id] = jobs_df.to_dict(orient='list')
    
    # 返回图表ID，即使这里没有实际的图表
    return chart_id

# 历史记录管理
def save_history(user_msg, assistant_msg, function_name=None, params=None, result=None, chart_id=None):
    if 'history' not in st.session_state:
        st.session_state.history = []
    
    history_item = {
        'user': user_msg,
        'assistant': assistant_msg,
        'timestamp': pd.Timestamp.now(),
        'function_name': function_name,
        'params': params,
        'result': result,
        'chart_id': chart_id
    }
    
    st.session_state.history.append(history_item)
    
    try:
        history_df = pd.DataFrame(st.session_state.history)
        history_df.to_csv('chat_history.csv', index=False)
    except Exception as e:
        st.warning(f"历史记录保存失败: {str(e)}")

def load_history():
    if 'history' not in st.session_state:
        try:
            if os.path.exists('chat_history.csv'):
                history_df = pd.read_csv('chat_history.csv')
                st.session_state.history = history_df.to_dict('records')
            else:
                st.session_state.history = []
        except Exception as e:
            st.session_state.history = []
            st.warning(f"历史记录加载失败: {str(e)}")
    return st.session_state.history

def show_history():
    history = load_history()
    
    if not history:
        st.info("暂无历史记录")
        return
    
    st.subheader("历史对话记录")
    
    col1, col2 = st.columns(2)
    with col1:
        search_text = st.text_input("搜索对话", "")
    with col2:
        max_results = st.slider("显示条数", 5, 50, 10)
    
    filtered_history = history
    if search_text:
        filtered_history = [h for h in history 
                           if search_text.lower() in h['user'].lower() 
                           or search_text.lower() in h['assistant'].lower()]
    
    if not filtered_history:
        st.info("未找到匹配的对话")
        return
    
    for i, item in enumerate(filtered_history[-max_results:]):
        with st.expander(f"对话 {len(filtered_history) - i} - {item['timestamp']}"):
            st.markdown(f"**用户**: {item['user']}")
            st.markdown(f"**助手**: {item['assistant']}")
            
            # 可视化历史记录
            try:
                if item['function_name'] == "recommend_related_skills" and item.get('params') and 'skill' in item['params'] and item.get('result'):
                    st.subheader("关联技术可视化")
                    chart_id = item.get('chart_id')
                    if chart_id and chart_id in st.session_state.charts:
                        st_echarts(options=st.session_state.charts[chart_id], key=chart_id)
                elif item['function_name'] == "query_skills_for_position" and item.get('params') and 'position' in item['params'] and item.get('result'):
                    st.subheader("岗位技能需求可视化")
                    chart_id = item.get('chart_id')
                    if chart_id and chart_id in st.session_state.charts:
                        st_echarts(options=st.session_state.charts[chart_id], key=chart_id)
                elif item['function_name'] == "recommend_positions_by_skill" and item.get('params') and 'skill' in item['params'] and item.get('result'):
                    st.subheader("技能适配岗位可视化")
                    chart_id = item.get('chart_id')
                    if chart_id and chart_id in st.session_state.charts:
                        st_echarts(options=st.session_state.charts[chart_id], key=chart_id)
                elif item['function_name'] == "query_highest_salary_jobs_by_city" and item.get('params') and 'city' in item['params'] and item.get('result'):
                    st.subheader("城市高薪岗位可视化")
                    chart_id = item.get('chart_id')
                    if chart_id and chart_id in st.session_state.charts:
                        st_echarts(options=st.session_state.charts[chart_id], key=chart_id)
                elif item['function_name'] == "query_jobs_by_education" and item.get('params') and 'education' in item['params'] and item.get('result'):
                    st.subheader("学历适配职业可视化")
                    chart_id = item.get('chart_id')
                    if chart_id and chart_id in st.session_state.charts:
                        jobs_df = pd.DataFrame(st.session_state.charts[chart_id])
                        st.dataframe(jobs_df, use_container_width=True)
            except Exception as e:
                st.warning(f"可视化失败: {str(e)}")
            
            col1, col2 = st.columns([1, 3])
            with col1:
                if st.button("重新运行", key=f"rerun_{i}"):
                    st.session_state.messages.append({"role": "user", "content": item['user']})
                    st.experimental_rerun()

# 主程序
# 主程序
def main(api_key="", temperature=0.6, messages=[]):
    with st.sidebar:
        st.title("系统设置")
        api_key = st.text_input("🔑 智谱AI API Key", type="password", value=api_key)
        st.session_state.final_api_key = api_key
        temperature = st.slider("🌡️ 回答温度", 0.1, 1.0, temperature, 0.1)
        st.session_state.final_temperature = temperature

        st.divider()
        st.subheader("问题功能示例")
        st.markdown("1.从技术问关联技术")
        st.markdown("例： 我擅长python，请推荐关联技术")
        st.markdown("2.从岗位问技术")
        st.markdown("例： 算法工程师需要学什么？")
        st.markdown("3.从技术问岗位")
        st.markdown("例子： 我会Java，什么岗位适合我")
        st.markdown("4.询问某地区最高薪资的岗位")
        st.markdown("例： 深圳最高薪资的前五个岗位是？")
        st.markdown("5.给学历，推荐岗位")
        st.markdown("例： 我本科生，推荐一些岗位")

        st.divider()
        st.subheader("历史记录")
        if st.button("查看历史对话"):
            st.session_state.show_history = not st.session_state.get('show_history', False)

        st.divider()
        if st.button("🚀 清空对话历史"):
            st.session_state.messages = [
                {"role": "assistant", "content": "你好！我是招聘智能助手，可以回答技术关联、岗位技能、岗位推荐等问题。"}
            ]
            st.success("对话历史已清空")

    if not api_key:
        st.warning("请在侧边栏输入智谱AI API Key")
        return

    if 'messages' not in st.session_state:
        st.session_state.messages = messages

    if st.session_state.get('show_history', False):
        show_history()
        return

    st.header("对话窗口")
    # 显示对话历史
    for msg in st.session_state.messages:
        with st.chat_message(msg['role']):
            st.markdown(msg['content'])
            # 检查是否有图表ID
            if "chart_id" in msg:
                chart_id = msg["chart_id"]
                if chart_id and chart_id in st.session_state.charts:
                    if isinstance(st.session_state.charts[chart_id], dict):
                        st_echarts(options=st.session_state.charts[chart_id], key=chart_id)
                    else:
                        jobs_df = pd.DataFrame(st.session_state.charts[chart_id])
                        st.dataframe(jobs_df, use_container_width=True)

    if prompt := st.chat_input("请输入问题"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            with st.spinner("正在分析..."):
                try:
                    client = init_ai(api_key)
                    data = load_and_preprocess_data()

                    if data.empty:
                        st.error("数据加载失败，请检查文件路径和格式")
                        return

                    frequent_itemsets, rules = generate_association_rules(data)

                    if rules.empty:
                        st.warning("警告：未生成有效关联规则，请检查数据或调整参数")
                        return

                    all_skills = {s for sublist in data['skill_list'] for s in sublist}
                    tools = [
                        recommend_related_skills_tool,
                        query_skills_for_position_tool,
                        recommend_positions_by_skill_tool,
                        query_highest_salary_jobs_by_city_tool,
                        query_jobs_by_education_tool
                    ]

                    response = chat_tools(client, prompt, tools, temperature)

                    if response.choices[0].message.tool_calls:
                        tool_call = response.choices[0].message.tool_calls[0]
                        function_name = tool_call.function.name
                        arguments = json.loads(tool_call.function.arguments)

                        try:
                            chart_id = None
                            if function_name == "recommend_related_skills":
                                skill_name = arguments['skill']
                                result = recommend_related_skills(
                                    skill_name,
                                    rules,
                                    all_skills
                                )

                                # 可视化
                                if not isinstance(result[0], str) or not result[0].startswith("数据集未收录"):
                                    st.subheader("关联技术可视化")
                                    chart_id = visualize_related_skills(skill_name, result)

                            elif function_name == "query_skills_for_position":
                                position_name = arguments['position']
                                result = query_skills_for_position(
                                    position_name,
                                    data
                                )

                                # 可视化
                                if not isinstance(result[0], str) or not result[0].startswith("数据集未收录"):
                                    st.subheader("岗位技能需求可视化")
                                    chart_id = visualize_skills_for_position(position_name, result)

                            elif function_name == "recommend_positions_by_skill":
                                skill_name = arguments['skill']
                                result = recommend_positions_by_skill(
                                    skill_name,
                                    data
                                )

                                # 可视化
                                if not isinstance(result[0], str) or not result[0].startswith("数据集未显示"):
                                    st.subheader("技能适配岗位可视化")
                                    chart_id = visualize_positions_by_skill(skill_name, result)

                            elif function_name == "query_highest_salary_jobs_by_city":
                                city_name = arguments['city']
                                result = query_highest_salary_jobs_by_city(city_name, data)

                                # 可视化
                                if not isinstance(result, str) and result:
                                    st.subheader("城市高薪岗位可视化")
                                    chart_id = visualize_highest_salary_jobs_by_city(city_name, result)

                            elif function_name == "query_jobs_by_education":
                                education = arguments['education']
                                result = query_jobs_by_education(education, data)

                                # 可视化
                                chart_id = visualize_jobs_by_education(education, result)

                            else:
                                result = [f"未知函数调用: {function_name}"]
                        except Exception as e:
                            result = [f"查询失败：{str(e)}"]

                        # 生成自然语言回答
                        try:
                            prompt_text = generate_prompt(
                                prompt, 
                                [r[0] for r in result] if isinstance(result[0], tuple) else result, 
                                function_name,
                                skill_name=arguments.get('skill') if function_name in ["recommend_related_skills", "recommend_positions_by_skill"] else None,
                                position=arguments.get('position') if function_name in ["query_skills_for_position"] else None,
                                city=arguments.get('city') if function_name in ["query_highest_salary_jobs_by_city"] else None,
                                education=arguments.get('education') if function_name == "query_jobs_by_education" else None
                            )

                            if not result or (isinstance(result, list) and len(result) > 0 and isinstance(result[0], str) and (result[0].startswith("数据集未收录") or result[0].startswith("没有找到"))):
                                fallback_text = "在私有数据中未查到相关信息，以下是基于行业普遍情况的回答："
                                st.warning(fallback_text)
                                internet_response = chat_tools(client, prompt_text)
                                response_text = internet_response.choices[0].message.content
                                st.markdown(response_text)
                            else:
                                final_response = chat_tools(client, prompt_text)
                                response_text = final_response.choices[0].message.content
                                st.markdown(response_text)

                            # 保存历史记录，包括图表ID
                            assistant_message = {"role": "assistant", "content": response_text, "chart_id": chart_id}
                            st.session_state.messages.append(assistant_message)
                            save_history(
                                prompt, 
                                response_text,
                                function_name=function_name,
                                params=arguments,
                                result=result,
                                chart_id=chart_id
                            )

                        except Exception as e:
                            st.error(f"生成回答失败：{str(e)}")

                    else:
                        response_text = response.choices[0].message.content
                        st.markdown(response_text)
                        st.session_state.messages.append({"role": "assistant", "content": response_text})
                        save_history(prompt, response_text)

                except Exception as e:
                    error_text = f"处理失败：{str(e)}"
                    st.error(error_text)
                    st.session_state.messages.append({"role": "assistant", "content": error_text})

    # 保存消息状态
    st.session_state.final_messages = st.session_state.messages

if __name__ == "__main__":
    main()